# =================================================
# === Divide sclice with max area into 2D image ===
# =================================================

# ! <<< Self
import pathlib
# ! >>>
from collections import Counter
import os
import SimpleITK as sitk
import argparse
import numpy as np
import os  # 遍历文件夹
import nibabel as nib  # nii格式一般都会用到这个包
import imageio  # 转换成图像
import pydicom
from PIL import Image
import sys

center = -500  # 肺部的窗宽窗位
width = 500


def create_dir(path):
    """ Create a directory. """
    try:
        if not os.path.exists(path):
            os.makedirs(path)
    except OSError:
        print(f"[Error in create dir]: Creating directory with name {path}")


def from_HU_to_gray(hu, slope, intercept):
    return int((hu - intercept) / slope)


def nii_to_image(max2Dslice, filepath):
    # filenames = os.listdir(max2Dslice,filepath)  # 读取nii文件夹

    # for f in filenames:
    #     开始读取nii文件
    # img_path = os.path.join(filepath, f)
    # img = nib.load(img_path)  # 读取nii
    # img_fdata = img.get_fdata()  # api 已完成转换，读出来的即为CT值
    # fname = f.replace('.nii', '')  # 去掉nii的后缀名
    # img_f_path = os.path.join(filepath, fname)
    # 创建nii对应的图像的文件夹
    # if not os.path.exists(img_f_path):
    #     os.mkdir(img_f_path)  # 新建文件夹

    # 转换成窗宽窗位
    min_value = (2 * center - width) / 2.0 + 0.5
    max_value = (2 * center + width) / 2.0 + 0.5
    dFactor = 255.0 / (max_value - min_value)

    # 开始转换为图像
    # (x, y, z) = img.shape
    # for i in range(z):  # z是图像的序列
    # silce = img_fdata[:, :, i]  # 选择哪个方向的切片都可以
    silce = max2Dslice.squeeze(0)
    silce = silce - min_value
    silce = np.trunc(silce * dFactor)
    silce[silce < 0.0] = 0
    silce[silce > 255.0] = 255  # 转换为窗位窗位之后的数据
    # maskimg_slice = maskimg_fdata[:, :, i]
    # finalimg=silce[int((x - 512) / 2):int((x - 512) / 2) + 512]
    finalimg = Image.fromarray(silce).convert('RGB')
    temp = './png_data/' + '{}.png'.format(filepath)
    finalimg.save(temp)
    # imageio.imwrite(temp,
    #                 finalimg)


def save_dcm_as_RGB_png(max2D, save_path, as_rgb, filename, savename):
    # dcm_file_list = [image_root +
    #                  f for f in os.listdir(image_root) if f.endswith('.Dcm')]
    # print(f'find {dcm_file_list.__len__()} dcm images')
    # if (dcm_file_list.__len__() == 0):
    #     print('[Error in dcm finding]: No dcm dataset')
    #     sys.exit(1)
    # ref = pydicom.dcmread(dcm_file_list[0])
    # dcmSize = (int(ref.Rows), int(ref.Columns), len(dcm_file_list))
    # print(f'dcm image size: {dcmSize}')
    # create_dir(f'{save_path}/images')
    # print(f'reading dcm to image...')
    # for i in range(0, dcm_file_list.__len__()):
    ds = max2D.squeeze(0)
    ds_pixel_arr = ds
    # ds_pixel_arr = np.where(ds_pixel_arr > from_HU_to_gray(max_bound, slope, intercept),
    #                         from_HU_to_gray(max_bound, slope, intercept), ds_pixel_arr)
    # ds_pixel_arr = np.where(ds_pixel_arr < from_HU_to_gray(min_bound, slope, intercept),
    #                         from_HU_to_gray(min_bound, slope, intercept), ds_pixel_arr)

    # ds_pixel_arr = ds_pixel_arr - \
    #     from_HU_to_gray(min_bound, slope, intercept)
    # ds_pixel_arr = ds_pixel_arr / \
    #     (from_HU_to_gray(max_bound, slope, intercept) -
    #      from_HU_to_gray(min_bound, slope, intercept))
    # ds_pixel_arr = ds_pixel_arr * 255
    # ds_pixel_arr = ds_pixel_arr.astype(int)

    mode = 'RGB' if as_rgb else 'I'
    image = Image.fromarray(ds_pixel_arr).convert(mode)
    file_name = filename.split('.')[0]
    if not pathlib.Path(save_path + os.sep + savename).exists():
        pathlib.Path(save_path + os.sep + savename).mkdir(parents=True, exist_ok=True)
    image.save(f'{save_path}/{savename}/{file_name}.png')
    # print('\rread No.{} dcm image [{}%]\t\t'.format(
    #     i + 1, format((i + 1) / dcm_file_list.__len__() * 100, '.2f')), end='\r')
    # print('\r\n')
    print('dcm read done')


def normalize(volume, a_min, a_max, b_min=None, b_max=None):
    """Normalize the volume"""
    volume[volume < a_min] = a_min
    volume[volume > a_max] = a_max
    volume = (volume - a_min) / (a_max - a_min)
    if b_min is not None and b_max is not None:
        volume = (volume * (b_max - b_min)) + b_min
    volume = volume.astype('float32')
    return volume


def crop_by_intensity(img, mask):
    """
    img: numpy array of 3D(x,y,z)
    """
    nozero = np.nonzero(mask)

    minz = nozero[0].min()
    miny = nozero[1].min()
    minx = nozero[2].min()

    maxz = nozero[0].max()
    maxy = nozero[1].max()
    maxx = nozero[2].max()

    zcount = nozero[0]
    tablez = Counter(zcount)
    z_max2Dslice = tablez.most_common(1)[0][0]
    zc_max2Dslice = tablez.most_common(1)[0][1]

    xoy = np.sum(mask, axis=0)
    yoz = np.sum(mask, axis=0)
    xoz = np.sum(mask, axis=1)

    ox = np.sum(xoy, axis=1)
    oy = np.sum(xoy, axis=0)
    oz = np.sum(xoz, axis=1)

    x1 = (ox != 0).argmax()
    x2 = (np.flipud(ox) != 0).argmax()
    x = ox.shape[0] - x1 - x2
    y1 = (oy != 0).argmax()
    y2 = (np.flipud(oy) != 0).argmax()
    y = oy.shape[0] - y1 - y2
    z1 = (oz != 0).argmax()
    z2 = (np.flipud(oz) != 0).argmax()
    z = oz.shape[0] - z1 - z2

    # new_img = img[x1:x1+x, y1:y1+y, z1:z1+z]
    new_img2 = img[z1:z1+z, y1:y1 + y, x1:x1+x]
    new_img = img[minz:maxz+1, miny:maxy+1, minx:maxx+1]
    max2Dslice = img[z_max2Dslice:z_max2Dslice + 1, miny:maxy + 1, minx:maxx + 1]
    max2Dslice_mask = mask[z_max2Dslice:z_max2Dslice + 1, miny:maxy + 1, minx:maxx + 1]
    return new_img, max2Dslice, max2Dslice_mask

## using simpleITK to load and save data.


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--png_save_as_rgb', default=True, help='save dcm as RGB png image')
    opt = parser.parse_args()

    # filepath = 'data1'
    # nii_to_image(filepath)

    file_root = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_classification'
    # filepath = file_root + os.sep + 'training/img'
    # filepath_roi = file_root + os.sep + 'training/mask'
    # png_save_path = file_root + os.sep + 'training_2d_png'
    filepath = file_root + os.sep + 'validation/img'
    filepath_roi = file_root + os.sep + 'validation/mask'
    png_save_path = file_root + os.sep + 'validation_2d_png'

    filepath_pathlib = pathlib.Path(filepath)
    filepath_roi_pathlib = pathlib.Path(filepath_roi)

    for file in filepath_pathlib.rglob('*.nii.gz'):
        filename = file.name
        itk_img = sitk.ReadImage(file.as_posix(), sitk.sitkFloat32)
        img = sitk.GetArrayFromImage(itk_img)
        print("img shape:", img.shape)

        # ! Fix bug: Need normalization
        img = normalize(img, a_min=0, a_max=1500, b_min=0, b_max=255)

        file_roi = filepath_roi_pathlib.joinpath(filename)
        itk_img_roi = sitk.ReadImage(file_roi.as_posix())
        img_roi = sitk.GetArrayFromImage(itk_img_roi)
        print("img_roi shape:", img_roi.shape)
        # img_data = np.where(img_roi == 1, img, 0)
        img_data2, max2Dslice, max2Dslice_mask = crop_by_intensity(img, img_roi)
        max2Dslice_z = np.where(max2Dslice_mask == 1, max2Dslice, 0)
        save_dcm_as_RGB_png(max2Dslice, png_save_path, opt.png_save_as_rgb, filename, 'images')
        save_dcm_as_RGB_png(max2Dslice_z, png_save_path, opt.png_save_as_rgb, filename, 'images2')
        pngname = filename.split('.')[0]
        maskname = filename.split('.')[0] + '_mask'
        savename = png_save_path+'/{}.png'.format(maskname)
        # nii_to_image(max2Dslice,pngname)
        max2Dslice_mask = max2Dslice_mask * 255
        mask_img = Image.fromarray(max2Dslice_mask.squeeze(0)).convert('L').convert('1')
        if not pathlib.Path(png_save_path + os.sep + 'masks').exists():
            pathlib.Path(png_save_path + os.sep + 'masks').mkdir(parents=True, exist_ok=True)
        mask_img.save(png_save_path + '/masks/{}.png'.format(maskname))
        # imageio.imwrite(savename,max2Dslice_mask.squeeze(0))
        ## save
        # out = sitk.GetImageFromArray(img_data2)

        # sitk.WriteImage(out, "roi5/"+filename_roi+'.nii.gz')
